import random
import math
import matplotlib.pyplot as plt
import copy

# 设置窗口标题
fig = plt.figure()
fig.canvas.manager.set_window_title('旅行商问题（NP问题）')

# 后续代码保持不变
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体字体
# 城市坐标：城市名,Y轴,X轴
city = [['A', 20, 20], ['B', 110, 30], ['C', 110, 60], ['D', 180, 80], ['E', 50, 130], ['F', 150, 180]]
path = {}  # 城市间距离
names = []  # 城市名称

# 三角函数计算距离
def dx(a, b):
    dx = abs(a[1] - b[1])
    dy = abs(a[0] - b[0])
    # 四舍五入取整
    dis = round(math.sqrt(dx ** 2 + dy ** 2))
    return dis

# 计算所有城市间的距离
def calculate_distances():
    for i in range(len(city)):
        name = city[i][0]
        names.append(name)  # 追加城市名称
        y = city[i][1]
        x = city[i][2]
        a = [y, x]  # 当前城市坐标
        for j in range(len(city)):
            # 排除自己
            if city[j][0] != name:
                # 航线名称
                key = name + '-' + city[j][0]
                # 其它城市坐标
                b = [city[j][1], city[j][2]]
                path[key] = dx(a, b)
    print('里程表', path)

# 绘制地图
def plot_map():
    w = 200
    h = 200
    gap = 20  # 地图边距
    for c in city:
        if c[1] > h:
            h = c[1] + gap
        if c[2] > w:
            w = c[2] + gap
    for c in city:
        name = c[0]
        y = c[1]
        x = c[2]
        # 绘制圆圈与文字
        plt.scatter(x, y, s=(gap / 2) ** 2, edgecolors='black', facecolors='none')
        plt.text(x, y, name, ha='center', va='center')

# 初始化种群（路程规划要求城市名不得重复）
def start(size):
    first = names.pop(0)  # 移动第一个城市名，到始发地中
    print(f'种群规模{size}，最终基因总路程最短')
    data = []
    # 随机初始物种
    for _ in range(size):
        gene = []
        arr = copy.copy(names)
        for _ in range(len(names)):
            # 随机一个城市名
            rand = random.randint(0, len(arr) - 1)
            name = arr[rand]
            gene.append(name)
            # 删除已经选择的城市，避免重复选择
            arr.remove(name)
        # 添加始发地
        gene.insert(0, first)
        gene.append(first)
        # 追加到种群中
        data.append(gene)
    print('初始化种群', data)
    return data, first

# 适应度函数
def ok(arr):
    # 组成航线
    sum_dis = 0
    for i in range(len(arr) - 1):
        name = arr[i] + '-' + arr[i + 1]
        sum_dis += path[name]
    return sum_dis

# 轮盘选择方法
def roulette_wheel_selection(data, num):
    fitness_values = [1 / ok(gene) for gene in data]
    total_fitness = sum(fitness_values)
    selection_probabilities = [fitness / total_fitness for fitness in fitness_values]
    selected = []
    for _ in range(num):
        r = random.random()
        cumulative_probability = 0
        for i in range(len(data)):
            cumulative_probability += selection_probabilities[i]
            if cumulative_probability > r:
                selected.append(data[i])
                break
    return selected

# 交叉操作
def crossover(parent1, parent2):
    len_names = len(parent1) - 2
    step = len_names // 2
    child = parent1.copy()
    for j in range(1, len_names + 1, step):
        if j + step > len_names:
            break
        if random.random() > 0.5:
            for s in range(step):
                child[j + s] = parent2[j + s]
    child = child[1:-1]
    child_one = []
    arr = [n for n in names if n not in child]
    for one in child:
        if one in child_one:
            rand = random.randint(0, len(arr) - 1)
            name = arr[rand]
            child_one.append(name)
            arr.remove(name)
        else:
            child_one.append(one)
    child_one.insert(0, parent1[0])
    child_one.append(parent1[-1])
    return child_one

# 变异操作
def mutate(gene):
    len_gene = len(gene) - 2
    n1 = random.randint(1, len_gene)
    n2 = random.randint(1, len_gene)
    if n1 != n2:
        gene[n1], gene[n2] = gene[n2], gene[n1]
    return gene

# 绘制路径
def plot_path(path):
    plt.clf()
    plot_map()
    for i in range(len(path) - 1):
        current_city = next((c for c in city if c[0] == path[i]), None)
        next_city = next((c for c in city if c[0] == path[i + 1]), None)
        if current_city and next_city:
            plt.plot([current_city[2], next_city[2]], [current_city[1], next_city[1]], 'r-')
    plt.title(f'目前最短路程: {ok(path)}')
    plt.draw()

# 开始遗传
def heredity(data, first, size, good, sex, max_iter):
    min_path = None
    for iter_count in range(max_iter):
        # 适应值排序
        data.sort(key=lambda x: ok(x))

        # 替换历史最小值
        if min_path is None or ok(data[0]) < ok(min_path):
            min_path = data[0]
            print(f'目前最优秀的基因是：{min_path}')

        print(f'目前最短路程是：{min_path} 大约：{ok(min_path)}')
        # 显示本次排序的数据
        for d in data:
            print(d, ok(d))

        print(f'还剩{max_iter - iter_count - 1}进化次机会')

        # 新路径规划
        news = []
        is_l = len(data)
        good_i = round(good * is_l)
        news.extend(data[:good_i])

        # 轮盘选择
        sex_i = round(sex * is_l)
        sexs = roulette_wheel_selection(data, sex_i)

        for _ in range(size - good_i):
            parent1, parent2 = random.sample(sexs, 2)
            child = crossover(parent1, parent2)
            child = mutate(child)
            news.append(child)

        data = news
        plt.pause(1)
        plot_path(min_path)

    return min_path

# 主程序
calculate_distances()
plot_map()
size = 10  # 种群规模
good = 0.01  # 精英率
sex = 0.85  # 交叉率
max_iter = 5  # 最大迭代次数
data, first = start(size)
final_path = heredity(data, first, size, good, sex, max_iter)
print('最终最短路程是：', final_path, '大约：', ok(final_path))
plt.show()